from stable_baselines3.common.torch_layers import BaseFeaturesExtractor
import torch.nn as nn

class TransformerFeatureExtractor(BaseFeaturesExtractor):
    def __init__(self, observation_space, seq_len, feature_dim, d_model=64, nhead=4, num_layers=2):
        super().__init__(observation_space, features_dim=d_model)  # ✅ 必须加这一句

        self.seq_len = seq_len
        self.feature_dim = feature_dim

        self.input_linear = nn.Linear(feature_dim, d_model)
        self.layernorm = nn.LayerNorm(d_model)

        encoder_layer = nn.TransformerEncoderLayer(d_model=d_model, nhead=nhead, batch_first=True)
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)

        nn.init.xavier_uniform_(self.input_linear.weight)

    def forward(self, observations):
        x = observations["features"]
        batch_size = x.shape[0]
        x = x.view(batch_size, self.seq_len, self.feature_dim)
        x = self.input_linear(x)
        x = self.layernorm(x)
        x = self.transformer(x)
        return x.mean(dim=1)
